from io import StringIO
import joblib
import uvicorn
from fastapi import FastAPI, File, UploadFile
from starlette.middleware.cors import CORSMiddleware
from entities.RequerstData import *
import pandas as pd
from entities.Response import Response
from data_conversion import *
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from model import health_model_create, salary_model_create

pd.set_option('display.max_columns', None)
app = FastAPI()

# 创建数据库连接
engine = create_engine("mysql+mysqlconnector://root:123456@localhost:3306/Tables")
Session = sessionmaker(bind=engine)
session = Session()
# 允许跨域请求的域名列表，如果要允许所有来源的跨域请求，可以将 allow_origins 设置为 ["*"]
# 如果只允许特定来源的跨域请求，可以将 allow_origins 设置为特定的域名列表，例如 ["http://localhost", "http://example.com"]
# 你也可以在 allow_origins_regex 中使用正则表达式来匹配特定的域名
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["GET", "POST", "PUT", "DELETE"],
    allow_headers=["*"],
)


@app.post("/healthpredict")
def data_predict(data: Data, response_model=Response):
    print(data)
    print(type(data))
    # 加载标准化对象，训练和预测数据都需要一致
    scaler = joblib.load("./model/health_model/health_scaler.pkl")
    # data请求数据格式化-->转换成DataFrame-->热编码
    X = health_data_conversion(data)
    X = scaler.transform(X)
    model = None
    if data.mode == "lg":
        model = joblib.load("./model/health_model/logistic_regression.pkl")
    if data.mode == "dt":
        model = joblib.load("./model/health_model/decision_tree.pkl")
    if data.mode == "knn":
        model = joblib.load("./model/health_model/KNN.pkl")
    if data.mode == "by":
        model = joblib.load("./model/health_model/Bayesian.pkl")
    if data.mode == "svn":
        model = joblib.load("./model/health_model/SVM.pkl")
    if data.mode == "rf":
        model = joblib.load("./model/health_model/random_forest.pkl")
    y_pred = model.predict(X)
    print(y_pred)
    print(type(y_pred))
    msg = '您健康有风险' if y_pred == 1 else '您很健康'
    return Response(data=int(y_pred[0]), msg=msg, code=200)


@app.post("/salarypredict")
def data_predict(data: SalaryData, response_model=Response):
    print(data)
    print(type(data))
    # 加载标准化对象，训练和预测数据都需要一致
    scaler = joblib.load("./model/salary_model/salary_scaler.pkl")
    X = salary_data_conversion(data)
    X = scaler.transform(X)
    model = None
    if data.mode == "lg":
        model = joblib.load("./model/salary_model/logistic_regression.pkl")
    if data.mode == "dt":
        model = joblib.load("./model/salary_model/decision_tree.pkl")
    if data.mode == "knn":
        model = joblib.load("./model/salary_model/KNN.pkl")
    if data.mode == "by":
        model = joblib.load("./model/salary_model/Bayesian.pkl")
    if data.mode == "svn":
        model = joblib.load("./model/salary_model/SVM.pkl")
    if data.mode == "rf":
        model = joblib.load("./model/salary_model/random_forest.pkl")
    y_pred = model.predict(X)
    print(y_pred)
    print(type(y_pred))
    msg = '您薪资预估>50k' if y_pred == 1 else '您薪资预估≤50k'
    return Response(data=int(y_pred[0]), msg=msg, code=200)


# 模型训练
@app.get('/train')
async def data_train(mode: str):
    #   根据data中的mode 进行有选择的预测
    if mode == "health":
        print(mode)
        health_model_create.train_and_evaluate_models()
    if mode == "salary":
        print(mode)
        salary_model_create.train_and_evaluate_salary_models()
    return Response(data=1, msg=f"{mode}所有模型更新成功", code=200)


if __name__ == '__main__':
    uvicorn.run(app, host="localhost", port=8008)
